{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nI learned BPE works by reading Karpathy's (https://github.com/karpathy/minbpe) so all credit to him!\\n\\nHow BPE works and is implemented, TLDR: \\n    - We \\n\""
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import List, Dict, Tuple\n",
    "\n",
    "'''\n",
    "I learned BPE works by reading Karpathy's lovely (https://github.com/karpathy/minbpe) so all credit to him! \n",
    "This implementation takes a similar approach, but is more bare-bones if you just want to understand how the fundamentals \n",
    "work, no frills. \n",
    "\n",
    "How BPE works and is implemented, TLDR: \n",
    "    - We start with raw bytes (0-255) as our base tokens and iteratively merge the most frequent adjacent pairs\n",
    "    - Key functions: merge() combines token pairs into new tokens, get_stats() tracks pair frequencies\n",
    "    - BPE class inherits from Tokenizer and learns merges during training, applies them greedily during encoding\n",
    "    - Simple but slow O(ML^2) encoding - we scan the sequence repeatedly looking for mergeable pairs\n",
    "    - All merges get stored in self.merges dict mapping (token1,token2)->new_token_id for easy lookup\n",
    "\n",
    "Key subtleties: \n",
    "    - Production BPE implementations would use a trie for faster lookup so it's O(ML) complexity\n",
    "    - self.vocab and self.merges are sufficient statistics for encode and decode. train() exists only to \n",
    "    populate these two for use at inference time. \n",
    "    - self.vocab is directly used to decode in a trivial way using lookup \n",
    "    - self.merges is used to repeat the merging process we used during training to construct our vocabulary \n",
    "    with new, unseen text, so its specifies our \"recipe\" to tokenize. \n",
    "    - It's crucial we merge in the same order as in training, hence why self.merges is ordered; \n",
    "    if not, we could be trying to merge (a, b) when there are no instances of a or b yet in the \n",
    "    token list we're constructing.\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge(ids: List[int], pair: Tuple[int, int], idx: int) -> List[int]: \n",
    "    # for instance, if ids is [1, 2, 3, 1, 2, 4]\n",
    "    # and pair is [1, 2] -> idx=7\n",
    "    # then we output [7, 3, 7, 4]\n",
    "    # this is the core BPE operation - it takes a sequence of token IDs and replaces all occurrences \n",
    "    # of a specific pair with a new token ID\n",
    "    n = len(ids)\n",
    "    new = []\n",
    "    i = 0\n",
    "    while i < n:\n",
    "        if i < n-1 and (ids[i], ids[i+1]) == pair: \n",
    "            new.append(idx)\n",
    "            i += 2\n",
    "        else: \n",
    "            new.append(ids[i])\n",
    "            i += 1\n",
    "    return new \n",
    "\n",
    "def get_stats(ids: List[int]) -> Dict[Tuple[int, int], int]:\n",
    "    # counts frequencies of adjacent token pairs in a sequence\n",
    "    # this is used to identify which pairs to merge next during training\n",
    "    n = len(ids)\n",
    "    freqs = {}\n",
    "    for i in range(n-1): \n",
    "        pair = (ids[i], ids[i+1])\n",
    "        freqs[pair] = freqs.get(pair, 0) + 1\n",
    "    return freqs\n",
    "\n",
    "class Tokenizer: \n",
    "    def __init__(self): \n",
    "        self.merges = {} # (int, int) -> int, these are our \"rules\" we learn during training and use during inference \n",
    "        self.pattern = \"\" # str  \n",
    "        self.vocab = self._build_vocab() # takes merges -> vocab using 256 bytes as inirt\n",
    "\n",
    "    def train(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def encode(self): \n",
    "        raise NotImplementedError\n",
    "\n",
    "    def decode(self): \n",
    "        raise NotImplementedError\n",
    "\n",
    "    def _build_vocab(self):\n",
    "        # constructs vocabulary mapping from token IDs to their byte sequences\n",
    "        # starts with base vocabulary of 256 bytes and builds up using merge rules\n",
    "        vocab = {idx: bytes([idx]) for idx in range(256)}\n",
    "        for (i1, i2), idx in self.merges:\n",
    "            vocab[idx] = vocab[i1] + vocab[i2]\n",
    "\n",
    "        return vocab \n",
    "        \n",
    "\n",
    "class BPE(Tokenizer): \n",
    "    def __init__(self): \n",
    "        super().__init__()\n",
    "\n",
    "    def train(self, text: List[str], vocab_size: int): # corpus Dict[Word: freq] -> populate merges and vocab \n",
    "        # use get_stats() and merge() update merges and vocab \n",
    "        merges, vocab = {}, self.vocab # (int, int) -> int and {idx:bytes} where former determines the latter\n",
    "        \n",
    "        BASE_VOCAB_SIZE = 256\n",
    "        assert vocab_size >= BASE_VOCAB_SIZE\n",
    "        num_merges = vocab_size-BASE_VOCAB_SIZE\n",
    "\n",
    "        # convert text to sequence of byte IDs - this is our starting point\n",
    "        ids = list(text.encode(\"utf-8\")) # list of ints representing chars \n",
    "\n",
    "        for i in range(num_merges): \n",
    "            # find the most popular pair using get_stats\n",
    "            freqs = get_stats(ids)\n",
    "            if not freqs:  # If no more pairs to merge\n",
    "                break\n",
    "            best_pair = max(freqs, key=freqs.get)            \n",
    "            new_idx = BASE_VOCAB_SIZE+i\n",
    "            # then merge using merge() to update merges \n",
    "            ids = merge(ids, best_pair, new_idx) \n",
    "            # update merges and vocab \n",
    "            merges[best_pair] = new_idx \n",
    "            vocab[new_idx] = vocab[best_pair[0]] + vocab[best_pair[1]]\n",
    "\n",
    "        self.merges, self.vocab = merges, vocab\n",
    "\n",
    "    # readme \n",
    "    def encode(self, text: str) -> List[int]: # use all merges\n",
    "        # Start with raw bytes\n",
    "        ids = list(text.encode(\"utf-8\"))\n",
    "        \n",
    "        # Apply merges iteratively until no more can be applied\n",
    "        # this is a greedy algorithm - we keep merging pairs until we can't anymore\n",
    "        # production implementations use tries for O(n) complexity instead of O(n^2)\n",
    "        while True:\n",
    "            # Try to find a mergeable pair\n",
    "            merged = False\n",
    "            for i in range(len(ids) - 1): # O(L)\n",
    "                pair = (ids[i], ids[i+1])\n",
    "                if pair in self.merges: \n",
    "                    # Apply the merge\n",
    "                    new_ids = ids[:i] + [self.merges[pair]] + ids[i+2:] # O(L)\n",
    "                    ids = new_ids\n",
    "                    merged = True\n",
    "                    break\n",
    "            \n",
    "            # If no merges were applied in this pass, we're done\n",
    "            if not merged:\n",
    "                break\n",
    "                \n",
    "        return ids # O(ML^2), so very slow for long sequences \n",
    "\n",
    "    def decode(self, ids: List[int]) -> str: \n",
    "        # decoding is simple - just look up each token ID in the vocabulary to get its bytes\n",
    "        # then concatenate and decode back to UTF-8\n",
    "        byte_strings = [self.vocab[i] for i in ids]\n",
    "        return b''.join(byte_strings).decode('utf-8')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original text: The quick brown fox jumps over the lazy dog.\n",
      "\n",
      "Encoded tokens: [296]\n",
      "Number of tokens: 1\n",
      "Decoded text: The quick brown fox jumps over the lazy dog.\n",
      "\n",
      "Decoding matches original: True\n",
      "\n",
      "Testing on new text: A quick brown dog jumps over the lazy fox.\n",
      "Encoded tokens: [65, 32, 113, 117, 105, 99, 107, 32, 98, 114, 111, 119, 110, 32, 100, 111, 103, 32, 106, 117, 109, 112, 115, 32, 111, 118, 101, 114, 32, 116, 257, 108, 97, 122, 121, 32, 102, 111, 120, 46]\n",
      "Number of tokens: 40\n",
      "Decoded text: A quick brown dog jumps over the lazy fox.\n",
      "Decoding matches original: True\n"
     ]
    }
   ],
   "source": [
    "# test the correctness of our tokenizer\n",
    "text = \"The quick brown fox jumps over the lazy dog.\"\n",
    "print(f\"Original text: {text}\")\n",
    "\n",
    "# create and train it\n",
    "tokenizer = BPE()\n",
    "tokenizer.train(text, vocab_size=300)\n",
    "\n",
    "# test the encode/decode functionality\n",
    "encoded = tokenizer.encode(text)\n",
    "decoded = tokenizer.decode(encoded)\n",
    "\n",
    "print(f\"\\nEncoded tokens: {encoded}\")\n",
    "print(f\"Number of tokens: {len(encoded)}\")\n",
    "print(f\"Decoded text: {decoded}\")\n",
    "print(f\"\\nDecoding matches original: {text == decoded}\")\n",
    "\n",
    "# now try tokenizing on unseen text we didn't train on, swap fox and dog\n",
    "new_text = \"A quick brown dog jumps over the lazy fox.\" \n",
    "print(f\"\\nTesting on new text: {new_text}\")\n",
    "encoded_new = tokenizer.encode(new_text)\n",
    "decoded_new = tokenizer.decode(encoded_new)\n",
    "\n",
    "# hurra, it works!\n",
    "print(f\"Encoded tokens: {encoded_new}\")\n",
    "print(f\"Number of tokens: {len(encoded_new)}\")\n",
    "print(f\"Decoded text: {decoded_new}\")\n",
    "print(f\"Decoding matches original: {new_text == decoded_new}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "envi",
   "language": "python",
   "name": "lingua_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
